import json
import os.path
from pathlib import Path


import yolo_predict as yd
import seal_recogniton as sr
from SealRecognitionProcessor import SealRecognitionProcessor
from YOLOPredictionProcessor import YOLOPredictionProcessor


def seal_detect(image_path, output_path="output"):
    # image_path = Path("/path/to/your/image_or_folder")
    #
    # if image_path.is_file():
    #     print(f"{image_path} 是一个文件")
    # elif image_path.is_dir():
    #     print(f"{image_path} 是一个文件夹")
    # else:
    #     print(f"{image_path} 不存在或既不是文件也不是文件夹")

    if not os.path.exists(output_path):
        os.mkdir(output_path)
    model_path ="seal_sign_model.pt"
    extract_folder = "cropped_seals"
    yd.process_yolo_predictions(image_path, model_path, extract_folder)
    result_map=sr.process_seal_recognition(extract_folder)
    with open(os.path.join(output_path,"seal_rec.json"), "w", encoding="utf-8") as f:
        json.dump(result_map, f, ensure_ascii=False, indent=4)
    print(result_map)

if __name__ == '__main__':
    # 输出路径
    output_path="output/seal_recognition"

    # 初始化YOLO+掩膜处理模块
    yolo_predict=YOLOPredictionProcessor(
        model_path="seal_sign_model.pt",
        extract_folder="cropped_seals"
    )

    # 初始化图章文字提取模块
    seal_recogniton=SealRecognitionProcessor(
        input_dir="cropped_seals",
        output_dir="output/seal_recognition",
        save_res_json=True
    )

    # YOLO检测
    results = yolo_predict.process_predictions(
        input_folder="seal_examples",
        save_dict_json=True
    )
    # 图章文字提取
    results = seal_recogniton.process_images()
    with open(os.path.join(output_path,"seal_rec.json"), "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=4)
    print(results)

